import argparse
from pprint import pprint

from tqdm import tqdm

from file_utils import MyPath
from constant import ImageType
from image_utils import shp2raster, read_image_as_array
from metrics import Evaluator


def filter_repeat(list_tif_or_shp_path: list):
    set_repeat = {}
    for path in list_tif_or_shp_path:
        file_name = MyPath(path).stem.strip('_栅格结果').strip('_矢量结果')
        if (file_name in set_repeat and str(path).endswith('.shp')) or str(path).endswith('.roi.shp'):
            continue
        else:
            set_repeat[file_name] = path
    return list(set_repeat.values())


def main():
    parser = argparse.ArgumentParser(description='计算结果评价指标')
    parser.add_argument('predict_path', type=str)
    parser.add_argument('label_path', type=str)
    parser.add_argument('--shp_field', type=str, nargs='+', default=['class_id', 'spec_id'])
    parser.add_argument('--image_path', type=str, default=None)
    parser.add_argument('--num_class', type=int, default=2)
    parser.add_argument('--output_path', type=str, default='.')
    args = parser.parse_args()
    predict_path = MyPath(args.predict_path)
    label_path = MyPath(args.label_path)
    if len(args.shp_field) == 1:
        shp_field_predict = shp_field_label = args.shp_field[0]
    elif len(args.shp_field) >= 2:
        shp_field_predict, shp_field_label = args.shp_field[:2]
    else:
        shp_field_predict, shp_field_label = ['class_id', 'spec_id']
    image_path = args.image_path
    num_class = args.num_class
    output_path = args.output_path
    # list_image_path = []
    if predict_path.is_file() and label_path.is_file():
        list_image_path = [((predict_path, label_path), image_path)]
    elif predict_path.is_dir() and label_path.is_dir():
        list_image_path = list(
            zip(filter_repeat(predict_path.iter_dir(ImageType.pattern() + ['*.shp'], recursive=True)),
                filter_repeat(label_path.iter_dir(ImageType.pattern() + ['*.shp'], recursive=True))))
        if not image_path:
            list_image_path_ori = [''] * len(list_image_path)
        elif MyPath(image_path).is_file():
            list_image_path_ori = [image_path] * len(list_image_path)
        elif MyPath(image_path).is_dir():
            list_image_path_ori = list(MyPath(image_path).iter_dir(ImageType.pattern(), recursive=True))[
                                  :len(list_image_path)]
        else:
            raise ValueError('--image_path参数设置错误！文件或者目录不存在！')
        list_image_path = list(zip(list_image_path, list_image_path_ori))
    else:
        raise ValueError('仅支持标签和结果同时为文件格式或者同时为目录！')
    metrics = Evaluator(num_class)
    for (predict_path, label_path), image_path in tqdm(list_image_path):
        if predict_path.suffix == '.shp' and label_path.suffix == '.shp':
            assert image_path, '标签和结果都为shp格式时，原始影像路径不能为空！'
            predict_path_tif = predict_path.replace_suffix('.tif')
            shp2raster(str(image_path), str(predict_path), str(predict_path_tif), shp_field_predict)
            label_path_tif = label_path.replace_suffix('.tif')
            shp2raster(str(image_path), str(label_path), str(label_path_tif), shp_field_label)
        elif predict_path.suffix == '.shp' and label_path.suffix != '.shp':
            predict_path_tif = predict_path.replace_suffix('.tif')
            shp2raster(str(label_path), str(predict_path), str(predict_path_tif), shp_field_predict)
            label_path_tif = label_path
        elif predict_path.suffix != '.shp' and label_path.suffix == '.shp':
            predict_path_tif = predict_path
            label_path_tif = label_path.replace_suffix('.tif')
            shp2raster(str(predict_path), str(label_path), str(label_path_tif), shp_field_label)
        else:
            predict_path_tif, label_path_tif = predict_path, label_path
        predict = read_image_as_array(predict_path_tif.path)
        label = read_image_as_array(label_path_tif.path)
        metrics.add_batch(predict, label)

    results = metrics.cal_metrics(metrics=['pre', 'rec', 'mpre', 'mrec'])
    pprint(results)
    if output_path:
        import json
        results = json.dumps(results, indent=4)
        output_path = MyPath(output_path)
        if not output_path.suffix:
            output_path.makedirs()
            with open(output_path.joinpath('result_metrics.json').path, 'w') as file:
                file.write(results)
        elif output_path.suffix == '.json':
            output_path.make_parent_dirs()
            with open(output_path.path, 'w') as file:
                file.write(results)
        else:
            output_path.write_text(results)


if __name__ == '__main__':
    main()
